#include <ddm/AtomicFib.h>
#include <folly/logging/xlog.h>
#include <openr/messaging/ReplicateQueue.h>

#include <iostream>

namespace ddm {
AtomicFib::AtomicFib(
    openr::messaging::RQueue<openr::DecisionRouteUpdate> routeUpdatesQueue,
    openr::messaging::ReplicateQueue<AtomicRoute>& atomicRouteUpdatesQueue)
    : atomicRouteUpdatesQueue_(atomicRouteUpdatesQueue) {
  atoms_.insert(0);
  atoms_.insert(1L << 32);

  // Fiber to process route updates
  addFiberTask([q = std::move(routeUpdatesQueue), this]() mutable noexcept {
    while (true) {
      auto maybeThriftObj = q.get();  // perform read
      if (maybeThriftObj.hasError()) {
        XLOG(DBG1) << "Terminating route delta processing fiber";
        break;
      }

      processRouteUpdate(std::move(maybeThriftObj).value());
      dump();
    }
  });
}

void AtomicFib::stop() { openr::OpenrEventBase::stop(); }

void AtomicFib::processRouteUpdate(openr::DecisionRouteUpdate&& routeUpdate) {
  std::cout << "size" << atoms_.size() << std::endl;
  // Add/Update unicast routes to update
  for (const auto& [prefix, route] : routeUpdate.unicastRoutesToUpdate) {
    insertRoute(prefix, route);
  }
  for (const auto& prefix : routeUpdate.unicastRoutesToDelete) {
    deleteRoute(prefix);
  }
  computeAtomicRouteUpdate();
}

void AtomicFib::insertRoute(folly::CIDRNetwork prefix,
                            openr::RibUnicastEntry route) {
  long low = prefix.first.asV4().mask(prefix.second).toLongHBO();
  long high = low + (1 << (32 - prefix.second));
  std::cout << low << " " << high << std::endl;

  if (atoms_.find(low) == atoms_.end()) {
    long below = *beupper_bound(low);
    for (auto& it : atomToRules_[below]) {
      atomToRules_[low].insert(it);
    }
    atoms_.insert(low);
  }
  if (atoms_.find(high) == atoms_.end()) {
    long below = *beupper_bound(high);
    for (auto& it : atomToRules_[below]) {
      atomToRules_[high].insert(it);
    }
    atoms_.insert(high);
  }

  std::set<long>::iterator it = atoms_.lower_bound(low);
  while (*it < high) {
    atomToRules_[*it][prefix.second] = route;
    ++it;
  }
}

void AtomicFib::deleteRoute(folly::CIDRNetwork prefix) {
  long low = prefix.first.asV4().mask(prefix.second).toLongHBO();
  long high = low + (1 << (32 - prefix.second));

  std::set<long>::iterator it = atoms_.lower_bound(low);
  while (*it < high) {
    atomToRules_[*it].erase(prefix.second);
    ++it;
  }
}

void AtomicFib::computeAtomicRouteUpdate() {
  std::map<std::string, std::set<long>> egressToAtoms;
  for (const auto& [atom, routes] : atomToRules_) {
    if (routes.size() == 0) {
      continue;
    }
    auto highestPriorityRoute = routes.rbegin()->second;
    // We use the first nexthop as egress, ECMP is not implemented
    std::string egress =
        highestPriorityRoute.nexthops.begin()->address()->ifName().value_or(
            "N/A");
    egressToAtoms[egress].insert(atom);
  }

  // TODO: capture the deletion of atomic routes
  std::vector<std::string> affectedEgresses;
  for (const auto& [egress, atoms] : egressToAtoms) {
    if (egressToAtoms_.count(egress) == 0) {
      affectedEgresses.push_back(egress);
      continue;
    }
    if (atoms != egressToAtoms_[egress]) {
      affectedEgresses.push_back(egress);
    }
  }

  // reset cache
  egressToAtoms_ = egressToAtoms;

  for (const auto& egress : affectedEgresses) {
    std::vector<std::pair<long, long>> ranges;
    for (const auto& atom : egressToAtoms[egress]) {
      long high = *(std::next(atoms_.find(atom)));
      ranges.push_back(std::make_pair(atom, high));
    }
    AtomicRoute atomicRoute(egress, ranges);
    atomicRouteUpdatesQueue_.push(std::move(atomicRoute));
  }
}

void AtomicFib::dump() {
  std::cout << "=== dumping AtomicFib ===" << std::endl;
  for (auto it : atoms_) {
    std::cout << it << std::endl;
    for (const auto& [mask, route] : atomToRules_[it]) {
      std::cout << " " << mask << ", "
                << (*route.nexthops.begin()).address()->ifName().value_or("N/A")
                << std::endl;
    }
  }
  std::cout << "=== dumping AtomicFib end ===" << std::endl;
}

std::set<long>::iterator AtomicFib::beupper_bound(long atom) {
  return --(atoms_.upper_bound(atom));
}

}  // namespace ddm